from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.base import TaskResult
from autogen_agentchat.conditions import TextMentionTermination, ExternalTermination, TextMessageTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient
model_client = OpenAIChatCompletionClient(model="qwen2.5:7b",
                                                 model_info={
                                                     "vision": False,
                                                     "function_calling": True,
                                                     "family": "Qwen3",
                                                     "structured_output": True,
                                                     "json_output": True,
                                                 },
                                                 api_key="ollama",
                                          base_url="http://127.0.0.1:11434/v1")


# Create a tool to add 1 to a number.
def increment_number(number: int) -> int:
    """Increment a number by 1."""
    return number + 1


# 创建一个使用 increment_number 函数的工具代理。
looped_assistant = AssistantAgent(
    "looped_assistant",
    model_client=model_client,
    tools=[increment_number],  # 注册工具。
    system_message="You are a helpful AI assistant, use the tool to increment the number.",
)





async def main():
    # 当代理以文本消息响应时停止任务的终止条件。
    termination_condition = TextMessageTermination("looped_assistant")

    # 创建一个由循环助手代理和终止条件组成的团队。
    team = RoundRobinGroupChat(
        [looped_assistant],
        termination_condition=termination_condition,
    )

    # 使用一个任务运行团队，并将消息打印到控制台。
    await Console(team.run_stream(task="Increment the number 5 to 10."))
    # 重置团队
    await team.reset()
    await Console(team.run_stream(task="Increment the number to 20."))
if __name__ == "__main__":
    import asyncio
    asyncio.run(main())